package com.wsw.aitools.doubao;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import com.wsw.aitools.utils.ImageUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

/**
 * 使用火山引擎SDK
 */
@Slf4j
@Service
public class DoubaoApiV2 {
    
    public static final String apiKey = "158f035c-1d3e-4b0d-bf1c-63edef2fe1fd";
    public static final String modelName = "ep-20241224182005-8vh2m";
    // baseUrl 处填写火山方舟推理接入点/智能体资源所在区域对应的访问域名，如不填默认访问华北 2 (北京)。火山方舟支持的地域和访问域名参见 https://www.volcengine.com/docs/82379/1302013
    public static final String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";
    
    // 豆包接口地址    
    public static final String url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions";
    
    private ArkService createArkService(){
        ArkService arkService = ArkService.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .build();
        
        return arkService;
    }
    
    
    public String askDoubao(String prompt, String fileUrl){
        
        ArkService arkService = createArkService();

 
        
        RestTemplate restTemplate = new RestTemplate();

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(apiKey);

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", modelName);

        Map<String, Object> textContent = new HashMap<>();
        textContent.put("type", "text");
        textContent.put("text", prompt);
        

        Map<String, Object> imageUrl1 = new HashMap<>();
        imageUrl1.put("type", "image_url");
        imageUrl1.put("image_url", new HashMap<String, String>() {{

            put("url", ImageUtils.encodeImage(fileUrl));
        }});

//        Map<String, Object> imageUrl2 = new HashMap<>();
//        imageUrl2.put("type", "image_url");
//        imageUrl2.put("image_url", new HashMap<String, String>() {{
//            put("url", "https://portal.volccdn.com/obj/volcfe/cloud-universal-doc/upload_a81e3cdd3e30617ecd524a132fdb2736.png");
//        }});

        Map<String, Object> message = new HashMap<>();
        message.put("role", "user");
        message.put("content", new Object[]{textContent, imageUrl1});


        

        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setRole(ChatMessageRole.USER);
        chatMessage.setContent(new Object[]{textContent, imageUrl1});


        ChatCompletionRequest chatCompletionRequest = new ChatCompletionRequest();
        chatCompletionRequest.setModel(modelName);
        chatCompletionRequest.setMessages(Arrays.asList(chatMessage));


        Long start = System.currentTimeMillis();
        ChatCompletionResult chatCompletionResult = arkService.createChatCompletion(chatCompletionRequest);
        
        Long end = System.currentTimeMillis();
        log.info("豆包大模型请求耗时: " + (end - start) + "ms");
        
        log.info("chatCompletionResult: " + chatCompletionResult);

        
        return chatCompletionResult.getChoices().get(0).getMessage().getContent().toString();
    }
    
    
}
